import requests
from PIL import Image
from io import BytesIO
import sys
import os

# 设置控制台输出编码为utf-8
sys.stdout.reconfigure(encoding='utf-8')

def generate_image(prompt, width=768, height=1024, seed=42, model='flux'):
    """
    从 pollinations.ai 生成并下载图片
    
    参数:
        prompt: 图片描述文本
        width: 图片宽度
        height: 图片高度
        seed: 随机种子
        model: 使用的模型
    """
    # 获取当前脚本所在的目录
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 限制提示词长度，避免URL过长
    original_length = len(prompt)
    if original_length > 500:
        print(f'警告：提示词过长({original_length}字符)，已截断至500字符')
        prompt = prompt[:500]
    
    # 构建 URL
    image_url = f"https://pollinations.ai/p/{prompt}?width={width}&height={height}&seed={seed}&model={model}"
    
    try:
        # 获取图片
        response = requests.get(image_url)
        response.raise_for_status()  # 检查是否成功
        
        # 将响应内容转换为图片
        image = Image.open(BytesIO(response.content))
        
        # 保存图片到脚本所在目录
        output_filename = os.path.join(script_dir, 'ai_image.jpg')
        image.save(output_filename)
        print(f'图片已成功生成并保存为: {output_filename}')
        
        # 将URL保存到脚本所在目录的文件中
        info_filename = os.path.join(script_dir, 'image_info.txt')
        with open(info_filename, 'w', encoding='utf-8') as f:
            f.write(f'图片URL: {image_url}\n')
            f.write(f'提示词: {prompt}\n')
        print(f'图片URL和提示词已保存到: {info_filename}')
        
        return image
    except Exception as e:
        print(f'生成图片时发生错误: {str(e)}')
        return None

# 从prompt.txt读取提示词
def read_prompt_from_file(filename='prompt.txt'):
    """
    从文件中读取提示词
    
    参数:
        filename: 提示词文件的名称，默认为 'prompt.txt'
    返回:
        成功时返回提示词文本，失败时返回None
    """
    try:
        # 获取当前脚本所在的目录
        script_dir = os.path.dirname(os.path.abspath(__file__))
        # 构建prompt.txt的完整路径
        prompt_path = os.path.join(script_dir, filename)
        
        with open(prompt_path, 'r', encoding='utf-8') as f:
            prompt = f.read().strip()
        if not prompt:
            print(f'警告：{filename} 文件内容为空')
            return None
        print(f'成功读取提示词文件: {prompt_path}')
        return prompt
    except FileNotFoundError:
        print(f'错误：找不到 {filename} 文件')
        return None
    except Exception as e:
        print(f'读取 {filename} 时发生错误: {str(e)}')
        return None

# 主函数
def main():
    # 生成测试图片
    prompt = read_prompt_from_file()
    if prompt:
        generated_image = generate_image(
            prompt=prompt,
            width=768,
            height=1024,
            seed=42
        )
    else:
        print('由于无法读取提示词，将使用默认提示词')
        prompt = 'Spring sunshine, layered mountains, emerald rivers, soft golden light, traditional Chinese landscape painting style.'
        generated_image = generate_image(
            prompt=prompt,
            width=768,
            height=1024,
            seed=42
        )

# 如果直接运行此脚本，则执行主函数
if __name__ == "__main__":
    main()
